import unittest
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.common.by import By
from selenium import webdriver
import time


class TestLogin(unittest.TestCase):
    def setUp(self):
        self.driver = webdriver.Edge()
        self.driver.implicitly_wait(10)
        self.wait = WebDriverWait(self.driver, 10)
        self.driver.get("https://www.baidu.com")

    def tearDown(self):
        self.driver.close()

    def test_login(self):
        input = self.wait.until(lambda x: x.find_element(By.ID, "kw"))
        input.send_keys("张三")
        botton = self.wait.until(lambda x: x.find_element(By.ID, "su"))
        botton.click()
        time.sleep(5)
        assert "张三" in self.driver.page_source

    def test_login2(self):
        input = self.wait.until(lambda x: x.find_element(By.ID, "kw"))
        input.send_keys("李四")
        botton = self.wait.until(lambda x: x.find_element(By.ID, "su"))
        botton.click()
        time.sleep(5)
        assert "李四" in self.driver.page_source


if __name__ == '__main__':
    unittest.main()
